# enhanced_image_generation.py - 使用经过验证的快速配置
import os
import json
import base64
import time
import uuid
import logging
import requests
from typing import Dict, Any, List, Optional
from pydantic import BaseModel
from dotenv import load_dotenv

load_dotenv()
logger = logging.getLogger(__name__)

# 配置
STABILITY_AI_KEY = os.getenv("TEXT_TO_IMAGE_API_KEY")
IMAGES_DIR = "static/images"
TEMP_DIR = "static/temp"

# 确保目录存在
for directory in [IMAGES_DIR, TEMP_DIR]:
    os.makedirs(directory, exist_ok=True)

# 经过验证的引擎配置 - 基于诊断结果
PROVEN_ENGINE = "stable-diffusion-v1-6"  # ✅ 5.7秒验证成功
BACKUP_ENGINE = "stable-diffusion-xl-1024-v1-0"  # 备用引擎


class ImageEditRequest(BaseModel):
    """图像编辑请求模型"""
    image_id: str
    edit_type: str
    parameters: Dict[str, Any]


class ImageGenerationRequest(BaseModel):
    """图像生成请求模型"""
    prompt: str
    style: str = "realistic"
    composition: str = "center"
    color_scheme: str = "natural"
    vase_style: str = "modern"
    lighting: str = "soft"
    angle: str = "front"


class AdvancedImageGenerator:
    """使用经过验证配置的图像生成器"""
    
    def __init__(self, api_key: str = STABILITY_AI_KEY):
        self.api_key = api_key

    async def generate_enhanced_image(self, request: ImageGenerationRequest) -> Dict[str, Any]:
        """生成图像 - 使用经过验证的快速配置"""
        try:
            if not self.api_key:
                return {"error": "请配置Stability AI API密钥"}

            # 构建提示词
            enhanced_prompt = self._build_prompt(request)
            
            logger.info(f"🎨 使用验证配置生成: {enhanced_prompt[:50]}...")

            # 首先使用经过验证的快速引擎
            result = await self._generate_with_proven_engine(enhanced_prompt)
            if result and "error" not in result:
                return result

            # 如果快速引擎失败，尝试备用引擎
            logger.info("🔄 快速引擎失败，尝试备用引擎...")
            result = await self._generate_with_backup_engine(enhanced_prompt)
            if result and "error" not in result:
                return result

            return {"error": "所有验证的引擎都失败"}

        except Exception as e:
            logger.error(f"图像生成异常: {str(e)}")
            return {"error": f"生成失败: {str(e)}"}

    def _build_prompt(self, request: ImageGenerationRequest) -> str:
        """构建优化的提示词"""
        # 基于request构建提示词
        components = [
            "A beautiful professional floral arrangement",
            request.prompt,
            f"{request.style} style",
            f"with {request.vase_style} vase",
            f"{request.lighting} lighting",
            "high quality, detailed"
        ]
        return ", ".join(components)

    async def _generate_with_proven_engine(self, prompt: str) -> Dict[str, Any]:
        """使用经过验证的快速引擎 (stable-diffusion-v1-6)"""
        try:
            headers = {
                "Content-Type": "application/json",
                "Accept": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }

            # 使用诊断中验证成功的参数
            payload = {
                "text_prompts": [
                    {"text": prompt, "weight": 1.0}
                ],
                "cfg_scale": 5.0,      # 与诊断测试相同
                "height": 512,         # 与诊断测试相同
                "width": 512,          # 与诊断测试相同
                "samples": 1,
                "steps": 20            # 与诊断测试相同
            }

            logger.info(f"🚀 使用经过验证的引擎: {PROVEN_ENGINE}")
            
            start_time = time.time()
            response = requests.post(
                f"https://api.stability.ai/v1/generation/{PROVEN_ENGINE}/text-to-image",
                headers=headers,
                json=payload,
                timeout=30  # 诊断显示5.7秒，30秒足够
            )
            elapsed_time = time.time() - start_time

            logger.info(f"📡 响应时间: {elapsed_time:.1f}秒")

            return await self._process_response(response, "proven_fast", elapsed_time)

        except requests.exceptions.Timeout:
            logger.warning("经过验证的引擎也超时了")
            return {"error": "快速引擎超时"}
        except Exception as e:
            logger.warning(f"快速引擎失败: {str(e)}")
            return {"error": f"快速引擎失败: {str(e)}"}

    async def _generate_with_backup_engine(self, prompt: str) -> Dict[str, Any]:
        """使用备用引擎"""
        try:
            headers = {
                "Content-Type": "application/json",
                "Accept": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }

            # 备用引擎使用高质量参数
            payload = {
                "text_prompts": [
                    {"text": prompt, "weight": 1.0}
                ],
                "cfg_scale": 7.0,
                "height": 1024,
                "width": 1024,
                "samples": 1,
                "steps": 30
            }

            logger.info(f"🎨 使用备用引擎: {BACKUP_ENGINE}")
            
            start_time = time.time()
            response = requests.post(
                f"https://api.stability.ai/v1/generation/{BACKUP_ENGINE}/text-to-image",
                headers=headers,
                json=payload,
                timeout=120  # 高质量模式需要更长时间
            )
            elapsed_time = time.time() - start_time

            logger.info(f"📡 备用引擎响应时间: {elapsed_time:.1f}秒")

            return await self._process_response(response, "backup_quality", elapsed_time)

        except requests.exceptions.Timeout:
            logger.warning("备用引擎也超时")
            return {"error": "备用引擎超时"}
        except Exception as e:
            logger.warning(f"备用引擎失败: {str(e)}")
            return {"error": f"备用引擎失败: {str(e)}"}

    async def _process_response(self, response, mode: str, elapsed_time: float) -> Dict[str, Any]:
        """处理API响应"""
        if response.status_code == 200:
            result = response.json()
            
            if "artifacts" in result and len(result["artifacts"]) > 0:
                image_data = result["artifacts"][0]["base64"]
                
                # 保存图像
                image_id = f"floral_{mode}_{int(time.time())}_{uuid.uuid4().hex[:8]}"
                image_filename = f"{image_id}.png"
                image_path = os.path.join(IMAGES_DIR, image_filename)

                with open(image_path, "wb") as f:
                    f.write(base64.b64decode(image_data))

                logger.info(f"✅ {mode}模式生成成功: {image_filename} ({elapsed_time:.1f}秒)")

                # 保存元数据
                metadata = {
                    "id": image_id,
                    "mode": mode,
                    "generation_time": elapsed_time,
                    "engine": PROVEN_ENGINE if mode.startswith("proven") else BACKUP_ENGINE,
                    "created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
                    "filename": image_filename
                }

                metadata_path = os.path.join(TEMP_DIR, f"{image_id}_metadata.json")
                with open(metadata_path, 'w', encoding='utf-8') as f:
                    json.dump(metadata, f, ensure_ascii=False, indent=2)

                return {
                    "image_id": image_id,
                    "main_image": f"/static/images/{image_filename}",
                    "variants": [],  # 简化，专注于基础生成
                    "metadata": metadata
                }
            else:
                return {"error": f"{mode}模式: API未返回图像"}
                
        elif response.status_code == 401:
            return {"error": "API密钥无效"}
        elif response.status_code == 402:
            return {"error": "账户余额不足"}
        elif response.status_code == 429:
            return {"error": "API调用频率过高，请稍后重试"}
        else:
            try:
                error_info = response.json()
                return {"error": f"{mode}模式API错误: {error_info}"}
            except:
                return {"error": f"{mode}模式HTTP错误: {response.status_code}"}


class InteractiveImageEditor:
    """简化的图像编辑器"""
    
    async def replace_vase(self, image_id: str, new_vase_style: str):
        return {"error": "图像编辑功能需要基础生成稳定后再启用"}
    
    async def change_flower_color(self, image_id: str, target_color: str):
        return {"error": "图像编辑功能需要基础生成稳定后再启用"}
    
    async def enhance_lighting(self, image_id: str, lighting_style: str):
        return {"error": "图像编辑功能需要基础生成稳定后再启用"}


# 创建实例
enhanced_generator = AdvancedImageGenerator()
interactive_editor = InteractiveImageEditor()


# 测试函数
def test_proven_generation():
    """测试经过验证的生成配置"""
    import asyncio
    
    async def test():
        print("🧪 测试经过验证的配置...")
        
        request = ImageGenerationRequest(
            prompt="红色玫瑰花束",
            style="realistic",
            vase_style="modern"
        )
        
        start_time = time.time()
        result = await enhanced_generator.generate_enhanced_image(request)
        elapsed_time = time.time() - start_time
        
        print(f"🎯 测试完成，用时: {elapsed_time:.1f}秒")
        print(f"📊 结果: {result}")
        
        if "error" not in result:
            print("✅ 配置验证成功！")
        else:
            print(f"❌ 配置验证失败: {result['error']}")
        
        return result
    
    return asyncio.run(test())


if __name__ == "__main__":
    test_proven_generation()
